// Copyright 2012 The golab Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cmd

import (
	"bytes"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"strconv"
	"time"
)

var (
	ErrCmdNotStarted = errors.New("Command never started")
)

type Result struct {
	Cmd        *Cmd          // the command that generated the exited process
	Stdout     []byte        // the stdout from the command
	Stderr     []byte        // the stderr from the command
	Err        string        // error returned from Run or Start
	Success    bool          // whether the program exited successfully
	StartTime  time.Time     // the start time of the exited process
	SystemTime time.Duration // the system CPU time of the exited process and its children
	UserTime   time.Duration // the user CPU time of the exited process and its children
}

// Save creates the following files if there is content to be placed in them.
//
//   <path>/<stepID>_<cmdID>.log
//   <path>/<stepID>_<cmdID>.out
//   <path>/<stepID>_<cmdID>.err
//
// Here the path should be <exp_id>/<date>/<host> as sent to
// os.MkdirAll(path, perm).  Having a strict requirement for structuring the
// data output allows standardization of post experiment processing scripts.
// The log file is always created (unless the experiment crashes), but the out
// and err files are only created if there is something to put in them.
func (r *Result) Save(path, stepID string) error {
	// Assume that path to Dir exists (parent should call MkdirAll).
	var f *os.File
	var err error
	f, err = os.Create(fmt.Sprintf("%s/%s_%s.log", path, stepID, r.Cmd.ID))
	if err != nil {
		e, ok := err.(*os.PathError)
		if ok {
			return e.Err
		}
		return err
	}
	fmt.Fprintf(f, "Success: %v\n", r.Success)
	fmt.Fprintf(f, "SystemTime: %v\n", r.SystemTime)
	fmt.Fprintf(f, "UserTime : %v\n", r.UserTime)
	fmt.Fprintf(f, "Error: %s\n", r.Err)
	err = f.Sync()
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}
	if len(r.Stdout) > 0 {
		f, err = os.Create(fmt.Sprintf("%s/%s_%s.out", path, stepID, r.Cmd.ID))
		if err != nil {
			e, ok := err.(*os.PathError)
			if ok {
				return e.Err
			}
			return err
		}
		_, err = f.Write(r.Stdout)
		if err != nil {
			return err
		}
		err = f.Sync()
		if err != nil {
			return err
		}
		err = f.Close()
		if err != nil {
			return err
		}
	}
	if len(r.Stderr) > 0 {
		f, err = os.Create(fmt.Sprintf("%s/%s_%s.err", path, stepID, r.Cmd.ID))
		if err != nil {
			e, ok := err.(*os.PathError)
			if ok {
				return e.Err
			}
			return err
		}
		_, err = f.Write(r.Stderr)
		if err != nil {
			return err
		}
		err = f.Sync()
		if err != nil {
			return err
		}
		err = f.Close()
	}
	return err
}

// A Process is an intermediate reference that allows Cmds to be reused.
type Process struct {
	Cmd *Cmd
	st  time.Time
	ec  *exec.Cmd
	o   *bytes.Buffer
	e   *bytes.Buffer
	err error
}

func (p *Process) finish(r *Result) error {
	if p.st.IsZero() {
		return ErrCmdNotStarted
	}
	r.Cmd = p.Cmd
	r.Stdout = p.o.Bytes()
	r.Stderr = p.e.Bytes()
	if p.ec != nil && p.ec.ProcessState != nil {
		r.Success = p.ec.ProcessState.Success()
		r.SystemTime = p.ec.ProcessState.SystemTime()
		r.UserTime = p.ec.ProcessState.UserTime()
	}
	if p.err != nil {
		r.Err = p.err.Error()
	}
	return p.err
}

// Wait will block until a subprocess exits.
func (p *Process) Wait(r *Result) error {
	if p.err == nil {
		p.err = p.ec.Wait()
	}
	return p.finish(r)
}

// Stop tries to kill the subprocess with SIGINT and then waits on the results.
// Note that we are ignoring ExitErrors because they are generated by the use
// of Process.Signal.
func (p *Process) Stop(r *Result) error {
	if p.err == nil {
		p.err = p.ec.Process.Signal(os.Interrupt)
		err := p.ec.Wait()
		_, ok := err.(*exec.ExitError)
		if !ok {
			p.err = err
		}
	}
	return p.finish(r)
}

// Kill tries to kill the subprocess with SIGKILL and then waits on the results.
// Note that we are ignoring ExitErrors because they are generated by the use
// of Process.Signal.
func (p *Process) Kill(r *Result) error {
	if p.err == nil {
		p.err = p.ec.Process.Signal(os.Kill)
		err := p.ec.Wait()
		_, ok := err.(*exec.ExitError)
		if !ok {
			p.err = err
		}
	}
	return p.finish(r)
}

type Cmd struct {
	ID   string   // Identification for the command
	Path string   // the path of the command to run
	Args []string // command line arguments
	Dir  string   // the working directory of the command
}

// SetID can be used if you don't want to think up a unique ID to use in a map.
func (c *Cmd) SetID() {
	c.ID = strconv.FormatInt(rand.Int63(), 10)
}

func (c *Cmd) setup() (*Process, error) {
	c.Args = append([]string{c.Path}, c.Args...)
	p := &Process{
		Cmd: c,
		o:   new(bytes.Buffer),
		e:   new(bytes.Buffer),
	}
	p.ec = &exec.Cmd{
		Args:   c.Args,
		Dir:    c.Dir,
		Stdout: p.o,
		Stderr: p.e,
	}
	p.ec.Path, p.err = exec.LookPath(c.Path)
	if p.err == nil {
		p.ec.Args[0] = p.ec.Path
	}
	return p, p.err
}

// Start creates an OS level subprocess, waits for it to end and returns the results.
func (c *Cmd) Run(r *Result) error {
	p, err := c.setup()
	if err != nil {
		return err
	}
	p.st = time.Now()
	p.err = p.ec.Run()
	return p.finish(r)
}

// Start creates an OS level subprocess and returns an intermediate reference.
func (c *Cmd) Start() (*Process, error) {
	p, err := c.setup()
	if err != nil {
		return p, err
	}
	p.st = time.Now()
	p.err = p.ec.Start()
	return p, p.err
}
